import time
import json
import os
from datetime import datetime
from contextlib import contextmanager

class PerformanceTimer:
    """
    性能计时器，用于记录各个环节的执行时间
    """
    def __init__(self):
        self.reset()
    
    def reset(self):
        """重置所有计时数据"""
        self.start_time = time.time()
        self.steps = {}
        self.api_calls = {}
        self.info = {}
        self.question = ""
        self.answer = ""
    
    def start_step(self, step_name):
        """开始记录一个步骤的时间"""
        self.steps[step_name] = {
            'start': time.time(),
            'end': None,
            'duration': None
        }
    
    def end_step(self, step_name):
        """结束记录一个步骤的时间"""
        if step_name in self.steps:
            self.steps[step_name]['end'] = time.time()
            self.steps[step_name]['duration'] = self.steps[step_name]['end'] - self.steps[step_name]['start']
    
    def start_api_call(self, api_name):
        """开始记录一个API调用的时间"""
        self.api_calls[api_name] = {
            'start': time.time(),
            'end': None,
            'duration': None
        }
    
    def end_api_call(self, api_name):
        """结束记录一个API调用的时间"""
        if api_name in self.api_calls:
            self.api_calls[api_name]['end'] = time.time()
            self.api_calls[api_name]['duration'] = self.api_calls[api_name]['end'] - self.api_calls[api_name]['start']
    
    def add_info(self, key, value):
        """添加额外信息"""
        self.info[key] = value
    
    def get_total_time(self):
        """获取总执行时间"""
        return time.time() - self.start_time
    
    def print_summary(self):
        """打印性能摘要"""
        print("\n===== 性能报告 =====")
        print(f"总执行时间: {self.get_total_time():.2f} 秒")
        
        print("\n步骤执行时间:")
        for step_name, step_data in sorted(self.steps.items(), key=lambda x: x[1]['duration'] or 0, reverse=True):
            duration = step_data['duration'] or 0
            percentage = (duration / self.get_total_time()) * 100 if self.get_total_time() > 0 else 0
            print(f"  {step_name}: {duration:.2f} 秒 ({percentage:.1f}%)")
        
        print("\nAPI调用时间:")
        for api_name, api_data in sorted(self.api_calls.items(), key=lambda x: x[1]['duration'] or 0, reverse=True):
            duration = api_data['duration'] or 0
            percentage = (duration / self.get_total_time()) * 100 if self.get_total_time() > 0 else 0
            print(f"  {api_name}: {duration:.2f} 秒 ({percentage:.1f}%)")
        
        if self.info:
            print("\n额外信息:")
            for key, value in self.info.items():
                print(f"  {key}: {value}")
        print("====================\n")
    
    def save_report(self, question, answer):
        """保存性能报告到文件"""
        self.question = question
        self.answer = answer
        
        # 创建报告目录
        report_dir = "performance_reports"
        if not os.path.exists(report_dir):
            os.makedirs(report_dir)
        
        # 生成报告文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
        filename = os.path.join(report_dir, f"report_{timestamp}.json")
        
        # 准备报告数据
        report_data = {
            'timestamp': datetime.now().isoformat(),
            'total_time': self.get_total_time(),
            'steps': {k: {'duration': v['duration']} for k, v in self.steps.items() if v['duration'] is not None},
            'api_calls': {k: {'duration': v['duration']} for k, v in self.api_calls.items() if v['duration'] is not None},
            'info': self.info,
            'question': question,
            'answer_preview': answer[:200] + "..." if len(answer) > 200 else answer
        }
        
        # 保存报告
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(report_data, f, ensure_ascii=False, indent=2)
        
        print(f"性能报告已保存到: {filename}")

class StepTimer:
    """
    步骤计时器上下文管理器
    """
    def __init__(self, step_name):
        self.step_name = step_name
    
    def __enter__(self):
        global_timer.start_step(self.step_name)
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        global_timer.end_step(self.step_name)
        return False

class APITimer:
    """
    API调用计时器上下文管理器
    """
    def __init__(self, api_name):
        self.api_name = api_name
    
    def __enter__(self):
        global_timer.start_api_call(self.api_name)
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        global_timer.end_api_call(self.api_name)
        return False

# 创建全局计时器实例
global_timer = PerformanceTimer()